refactor(sampler): consolidate sampling interface, part 1 #5312
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
Codecov Report
@@ Coverage Diff @@
## master #5312 +/- ##
==========================================
+ Coverage 83.33% 83.38% +0.04%
==========================================
Files 338 342 +4
Lines 18678 18745 +67
==========================================
+ Hits 15565 15630 +65
- Misses 3113 3115 +2
📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
Padarn
reviewed
Sep 2, 2022
torch_geometric/sampler/base.py
Outdated
| @abstractmethod | ||
| def __init__( | ||
| self, | ||
| data: Union[Data, HeteroData, Tuple[FeatureStore, GraphStore]], |
There was a problem hiding this comment.
These are going to get painful to continue to support.. do we have "data to graphstore"? Can we assume that you will always convert to the new representation? Maybe that doesn't make sense...
There was a problem hiding this comment.
Yeah, this is a great point. We do plan to consolidate behind Tuple[FeatureStore, GraphStore] in the near future, but probably in a separate PR.
2 tasks
mananshah99
changed the title
GraphStore
mananshah99
changed the title
rusty1s
approved these changes
Sep 7, 2022
torch_geometric/sampler/base.py
Outdated
| # * edge: a tensor of the indices in the original graph; e.g. to be used to | ||
| # obtain edge attributes. | ||
| class SamplerOutput(NamedTuple): | ||
| node: Union[torch.Tensor, Dict[NodeType, torch.Tensor]] |
There was a problem hiding this comment.
We might want to add an optional batch attribute which assigns each node to an example. This will be necessary to integrate with the new pyg-lib implementation.
There was a problem hiding this comment.
Done; although this may become more complicated once we try to simplify 'metadata'. Can revert later if necessary.
wsad1
reviewed
Sep 7, 2022
wsad1
approved these changes
Sep 7, 2022
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
…ric into remote_backend_2
mananshah99
added a commit
that referenced
this pull request
Sep 8, 2022
… 2 (#5365) This PR begins the effort to consolidate PyG's sampling interface in preparation for moving `sample(...)` behind the `GraphStore` interface. This effort is somewhat large in scope and will be broken into multiple PRs for ease of review. It resolves some TODOs from #5312, and introduces others that will be resolved in follow-up PRs. The major change removes `LinkNeighborSampler`, consolidating it under a the common `BaseSampler` interface with `NeighborSampler`. In doing so, it modifies the `BaseSampler` interface to support `sample_from_nodes` and `sample_from_edges`, since oftentimes edge-based sampling applies some transformations and then re-uses the same logic as node-based sampling. A sampler need not define both node-based and edge-based sampling, but if it doesn't, an appropriate error will be raised. Resolved TODOs: * `LinkNeighborSampler` and `NeighborSampler` interfaces are aligned * No more special handling of `edge_type_to_str` in `filter_*` * No more special handling of `perm_dict`, which removes the need for `edge_type_to_str` entirely
14 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR begins the effort to consolidate PyG's sampling interface in preparation for moving
sample(...)behind theGraphStoreinterface. This effort is somewhat large in scope and will be broken into multiple PRs for ease of review. There are many TODOs left in this PR; all will be resolved in forthcoming PRs (the code will become cleaner in short order).This first change moves
NeighborSamplerbehind a common base class with clearly defined inputs and outputs. Many more changes will be required before this base class defines a stable interface (e.g.LinkNeighborSamplerhas different semantics,NeighborSampler.perm_dictis used inNeighborLoader, etc.). But for now, it will act as a first step towards this goal.It further refactors
loader/utils.pyinto those that are required for sampling (moved intosampler/utils.py) and those that are required for filtering (remain inloader/utils.py).Note about integration with GraphStore. Moving sampling behind the graph store requires defining a clear interface for sampling across PyG. Since a graph store should not inherently be coupled with a sampling method used (e.g. one can imagine using a NeighborLoader and LinkNeighborLoader on the same graph store), we must define this interface independently from the GraphStore. One can then imagine the GraphStore defining supported samplers for different loaders (along with the defaults, which sample in-memory), and those samplers are then seamlessly used in the loaders.